#!/usr/bin/env python3
"""
AIQuant数据库初始化脚本

这个脚本用于初始化AIQuant系统的数据库，包括：
- 创建数据库表结构
- 初始化基础数据
- 设置TimescaleDB扩展
- 创建索引和约束
"""

import os
import sys
import asyncio
import argparse
import logging
from pathlib import Path
from typing import Optional
import asyncpg
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from aiquant.config.settings import get_settings

logger = logging.getLogger(__name__)

class DatabaseInitializer:
    """数据库初始化器"""
    
    def __init__(self):
        self.settings = get_settings()
        
    def create_database_if_not_exists(self):
        """创建数据库（如果不存在）"""
        print(f"检查数据库 {self.settings.db_name} 是否存在...")
        
        # 连接到postgres数据库来创建目标数据库
        conn = psycopg2.connect(
            host=self.settings.db_host,
            port=self.settings.db_port,
            user=self.settings.db_user,
            password=self.settings.db_password,
            database='postgres'
        )
        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        
        cursor = conn.cursor()
        
        try:
            # 检查数据库是否存在
            cursor.execute(
                "SELECT 1 FROM pg_catalog.pg_database WHERE datname = %s",
                (self.settings.db_name,)
            )
            
            if cursor.fetchone():
                print(f"✅ 数据库 {self.settings.db_name} 已存在")
            else:
                # 创建数据库
                cursor.execute(f'CREATE DATABASE "{self.settings.db_name}"')
                print(f"✅ 数据库 {self.settings.db_name} 创建成功")
                
        except Exception as e:
            print(f"❌ 创建数据库失败: {e}")
            raise
        finally:
            cursor.close()
            conn.close()
            
    async def setup_timescaledb(self):
        """设置TimescaleDB扩展"""
        print("设置TimescaleDB扩展...")
        
        conn = await asyncpg.connect(
            host=self.settings.db_host,
            port=self.settings.db_port,
            user=self.settings.db_user,
            password=self.settings.db_password,
            database=self.settings.db_name
        )
        
        try:
            # 创建TimescaleDB扩展
            await conn.execute("CREATE EXTENSION IF NOT EXISTS timescaledb CASCADE;")
            print("✅ TimescaleDB扩展设置完成")
            
        except Exception as e:
            print(f"⚠️ TimescaleDB扩展设置失败: {e}")
            print("系统将继续使用标准PostgreSQL功能")
        finally:
            await conn.close()
            
    async def create_tables(self):
        """创建数据库表"""
        print("创建数据库表...")
        
        conn = await asyncpg.connect(
            host=self.settings.db_host,
            port=self.settings.db_port,
            user=self.settings.db_user,
            password=self.settings.db_password,
            database=self.settings.db_name
        )
        
        try:
            # 用户表
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    id SERIAL PRIMARY KEY,
                    username VARCHAR(50) UNIQUE NOT NULL,
                    email VARCHAR(100) UNIQUE NOT NULL,
                    password_hash VARCHAR(255) NOT NULL,
                    is_active BOOLEAN DEFAULT TRUE,
                    is_admin BOOLEAN DEFAULT FALSE,
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                );
            """)
            
            # 策略表
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS strategies (
                    id SERIAL PRIMARY KEY,
                    name VARCHAR(100) NOT NULL,
                    description TEXT,
                    strategy_type VARCHAR(50) NOT NULL,
                    parameters JSONB,
                    is_active BOOLEAN DEFAULT TRUE,
                    created_by INTEGER REFERENCES users(id),
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                );
            """)
            
            # 投资组合表
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS portfolios (
                    id SERIAL PRIMARY KEY,
                    name VARCHAR(100) NOT NULL,
                    description TEXT,
                    initial_capital DECIMAL(20, 8) NOT NULL,
                    current_value DECIMAL(20, 8) DEFAULT 0,
                    cash_balance DECIMAL(20, 8) DEFAULT 0,
                    is_active BOOLEAN DEFAULT TRUE,
                    created_by INTEGER REFERENCES users(id),
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                );
            """)
            
            # 持仓表
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS positions (
                    id SERIAL PRIMARY KEY,
                    portfolio_id INTEGER REFERENCES portfolios(id),
                    symbol VARCHAR(20) NOT NULL,
                    quantity DECIMAL(20, 8) NOT NULL,
                    average_price DECIMAL(20, 8) NOT NULL,
                    current_price DECIMAL(20, 8),
                    market_value DECIMAL(20, 8),
                    unrealized_pnl DECIMAL(20, 8),
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                );
            """)
            
            # 订单表
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS orders (
                    id SERIAL PRIMARY KEY,
                    portfolio_id INTEGER REFERENCES portfolios(id),
                    strategy_id INTEGER REFERENCES strategies(id),
                    symbol VARCHAR(20) NOT NULL,
                    order_type VARCHAR(20) NOT NULL,
                    side VARCHAR(10) NOT NULL,
                    quantity DECIMAL(20, 8) NOT NULL,
                    price DECIMAL(20, 8),
                    status VARCHAR(20) DEFAULT 'pending',
                    filled_quantity DECIMAL(20, 8) DEFAULT 0,
                    filled_price DECIMAL(20, 8),
                    commission DECIMAL(20, 8) DEFAULT 0,
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                );
            """)
            
            # 交易记录表
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS trades (
                    id SERIAL PRIMARY KEY,
                    order_id INTEGER REFERENCES orders(id),
                    portfolio_id INTEGER REFERENCES portfolios(id),
                    symbol VARCHAR(20) NOT NULL,
                    side VARCHAR(10) NOT NULL,
                    quantity DECIMAL(20, 8) NOT NULL,
                    price DECIMAL(20, 8) NOT NULL,
                    commission DECIMAL(20, 8) DEFAULT 0,
                    realized_pnl DECIMAL(20, 8),
                    trade_time TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                );
            """)
            
            # 市场数据表（时间序列）
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS market_data (
                    time TIMESTAMP WITH TIME ZONE NOT NULL,
                    symbol VARCHAR(20) NOT NULL,
                    open_price DECIMAL(20, 8),
                    high_price DECIMAL(20, 8),
                    low_price DECIMAL(20, 8),
                    close_price DECIMAL(20, 8),
                    volume DECIMAL(20, 8),
                    amount DECIMAL(20, 8),
                    PRIMARY KEY (time, symbol)
                );
            """)
            
            # 尝试创建超表（TimescaleDB）
            try:
                await conn.execute("""
                    SELECT create_hypertable('market_data', 'time', 
                                           chunk_time_interval => INTERVAL '1 day',
                                           if_not_exists => TRUE);
                """)
                print("✅ 市场数据超表创建成功")
            except Exception as e:
                print(f"⚠️ 超表创建失败，使用普通表: {e}")
            
            # 性能指标表
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS performance_metrics (
                    id SERIAL PRIMARY KEY,
                    portfolio_id INTEGER REFERENCES portfolios(id),
                    date DATE NOT NULL,
                    total_return DECIMAL(10, 6),
                    daily_return DECIMAL(10, 6),
                    cumulative_return DECIMAL(10, 6),
                    volatility DECIMAL(10, 6),
                    sharpe_ratio DECIMAL(10, 6),
                    max_drawdown DECIMAL(10, 6),
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                    UNIQUE(portfolio_id, date)
                );
            """)
            
            # 风险指标表
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS risk_metrics (
                    id SERIAL PRIMARY KEY,
                    portfolio_id INTEGER REFERENCES portfolios(id),
                    date DATE NOT NULL,
                    var_95 DECIMAL(20, 8),
                    var_99 DECIMAL(20, 8),
                    expected_shortfall DECIMAL(20, 8),
                    beta DECIMAL(10, 6),
                    correlation DECIMAL(10, 6),
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                    UNIQUE(portfolio_id, date)
                );
            """)
            
            # 系统日志表
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS system_logs (
                    id SERIAL PRIMARY KEY,
                    level VARCHAR(20) NOT NULL,
                    logger VARCHAR(100) NOT NULL,
                    message TEXT NOT NULL,
                    module VARCHAR(100),
                    function VARCHAR(100),
                    line_number INTEGER,
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                );
            """)
            
            print("✅ 数据库表创建完成")
            
        except Exception as e:
            print(f"❌ 创建表失败: {e}")
            raise
        finally:
            await conn.close()
            
    async def create_indexes(self):
        """创建索引"""
        print("创建数据库索引...")
        
        conn = await asyncpg.connect(
            host=self.settings.db_host,
            port=self.settings.db_port,
            user=self.settings.db_user,
            password=self.settings.db_password,
            database=self.settings.db_name
        )
        
        try:
            indexes = [
                "CREATE INDEX IF NOT EXISTS idx_users_username ON users(username);",
                "CREATE INDEX IF NOT EXISTS idx_users_email ON users(email);",
                "CREATE INDEX IF NOT EXISTS idx_strategies_type ON strategies(strategy_type);",
                "CREATE INDEX IF NOT EXISTS idx_portfolios_created_by ON portfolios(created_by);",
                "CREATE INDEX IF NOT EXISTS idx_positions_portfolio_id ON positions(portfolio_id);",
                "CREATE INDEX IF NOT EXISTS idx_positions_symbol ON positions(symbol);",
                "CREATE INDEX IF NOT EXISTS idx_orders_portfolio_id ON orders(portfolio_id);",
                "CREATE INDEX IF NOT EXISTS idx_orders_symbol ON orders(symbol);",
                "CREATE INDEX IF NOT EXISTS idx_orders_status ON orders(status);",
                "CREATE INDEX IF NOT EXISTS idx_trades_portfolio_id ON trades(portfolio_id);",
                "CREATE INDEX IF NOT EXISTS idx_trades_symbol ON trades(symbol);",
                "CREATE INDEX IF NOT EXISTS idx_market_data_symbol ON market_data(symbol);",
                "CREATE INDEX IF NOT EXISTS idx_performance_metrics_portfolio_date ON performance_metrics(portfolio_id, date);",
                "CREATE INDEX IF NOT EXISTS idx_risk_metrics_portfolio_date ON risk_metrics(portfolio_id, date);",
                "CREATE INDEX IF NOT EXISTS idx_system_logs_level ON system_logs(level);",
                "CREATE INDEX IF NOT EXISTS idx_system_logs_created_at ON system_logs(created_at);",
            ]
            
            for index_sql in indexes:
                await conn.execute(index_sql)
                
            print("✅ 数据库索引创建完成")
            
        except Exception as e:
            print(f"❌ 创建索引失败: {e}")
            raise
        finally:
            await conn.close()
            
    async def insert_initial_data(self):
        """插入初始数据"""
        print("插入初始数据...")
        
        conn = await asyncpg.connect(
            host=self.settings.db_host,
            port=self.settings.db_port,
            user=self.settings.db_user,
            password=self.settings.db_password,
            database=self.settings.db_name
        )
        
        try:
            # 创建默认管理员用户
            from passlib.context import CryptContext
            pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
            
            admin_password = pwd_context.hash("admin123")
            
            await conn.execute("""
                INSERT INTO users (username, email, password_hash, is_admin)
                VALUES ($1, $2, $3, $4)
                ON CONFLICT (username) DO NOTHING;
            """, "admin", "admin@aiquant.com", admin_password, True)
            
            # 创建示例策略
            await conn.execute("""
                INSERT INTO strategies (name, description, strategy_type, parameters, created_by)
                VALUES ($1, $2, $3, $4, $5)
                ON CONFLICT DO NOTHING;
            """, "均值回归策略", "基于统计套利的均值回归策略", "mean_reversion", 
            '{"lookback_period": 20, "threshold": 2.0}', 1)
            
            await conn.execute("""
                INSERT INTO strategies (name, description, strategy_type, parameters, created_by)
                VALUES ($1, $2, $3, $4, $5)
                ON CONFLICT DO NOTHING;
            """, "动量策略", "基于价格动量的趋势跟踪策略", "momentum", 
            '{"fast_period": 10, "slow_period": 30}', 1)
            
            # 创建示例投资组合
            await conn.execute("""
                INSERT INTO portfolios (name, description, initial_capital, current_value, cash_balance, created_by)
                VALUES ($1, $2, $3, $4, $5, $6)
                ON CONFLICT DO NOTHING;
            """, "默认投资组合", "系统默认投资组合", 1000000.0, 1000000.0, 1000000.0, 1)
            
            print("✅ 初始数据插入完成")
            
        except Exception as e:
            print(f"❌ 插入初始数据失败: {e}")
            raise
        finally:
            await conn.close()
            
    async def verify_installation(self):
        """验证安装"""
        print("验证数据库安装...")
        
        conn = await asyncpg.connect(
            host=self.settings.db_host,
            port=self.settings.db_port,
            user=self.settings.db_user,
            password=self.settings.db_password,
            database=self.settings.db_name
        )
        
        try:
            # 检查表是否存在
            tables = await conn.fetch("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public'
                ORDER BY table_name;
            """)
            
            expected_tables = {
                'users', 'strategies', 'portfolios', 'positions', 
                'orders', 'trades', 'market_data', 'performance_metrics',
                'risk_metrics', 'system_logs'
            }
            
            existing_tables = {row['table_name'] for row in tables}
            missing_tables = expected_tables - existing_tables
            
            if missing_tables:
                print(f"❌ 缺少表: {missing_tables}")
                return False
            else:
                print(f"✅ 所有表都存在: {existing_tables}")
                
            # 检查数据
            user_count = await conn.fetchval("SELECT COUNT(*) FROM users;")
            strategy_count = await conn.fetchval("SELECT COUNT(*) FROM strategies;")
            portfolio_count = await conn.fetchval("SELECT COUNT(*) FROM portfolios;")
            
            print(f"✅ 数据统计: 用户({user_count}), 策略({strategy_count}), 投资组合({portfolio_count})")
            
            return True
            
        except Exception as e:
            print(f"❌ 验证失败: {e}")
            return False
        finally:
            await conn.close()


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="AIQuant数据库初始化脚本")
    parser.add_argument(
        "--force", "-f",
        action="store_true",
        help="强制重新初始化（危险操作）"
    )
    parser.add_argument(
        "--skip-timescale",
        action="store_true",
        help="跳过TimescaleDB设置"
    )
    parser.add_argument(
        "--verify-only",
        action="store_true",
        help="仅验证数据库状态"
    )
    
    args = parser.parse_args()
    
    # 设置日志
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    
    initializer = DatabaseInitializer()
    
    try:
        if args.verify_only:
            success = await initializer.verify_installation()
            sys.exit(0 if success else 1)
            
        print("🚀 开始初始化AIQuant数据库...")
        
        # 创建数据库
        initializer.create_database_if_not_exists()
        
        # 设置TimescaleDB
        if not args.skip_timescale:
            await initializer.setup_timescaledb()
            
        # 创建表
        await initializer.create_tables()
        
        # 创建索引
        await initializer.create_indexes()
        
        # 插入初始数据
        await initializer.insert_initial_data()
        
        # 验证安装
        success = await initializer.verify_installation()
        
        if success:
            print("\n🎉 数据库初始化完成！")
            print("\n默认管理员账户:")
            print("  用户名: admin")
            print("  密码: admin123")
            print("  邮箱: admin@aiquant.com")
            print("\n⚠️ 请在生产环境中修改默认密码！")
        else:
            print("\n❌ 数据库初始化失败！")
            sys.exit(1)
            
    except Exception as e:
        print(f"\n❌ 初始化过程中出错: {e}")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())